#include "websocket.h"
#include "file.h"
#include "fstream.h"
#include "http_server.h"
#include "stream.h"
#include "task.h"
#include "uring_loop.h"
#include <cstdint>
#include <cstdio>
#include <iostream>
#include <openssl/evp.h>
#include <openssl/sha.h>
#include <string>
#include <zlib.h>

namespace co_async {

uint64_t htonll(uint64_t hostlong) {
  uint64_t netlong = 0;
  const uint8_t *src = reinterpret_cast<const uint8_t *>(&hostlong);
  uint8_t *dest = reinterpret_cast<uint8_t *>(&netlong);

  dest[0] = src[7];
  dest[1] = src[6];
  dest[2] = src[5];
  dest[3] = src[4];
  dest[4] = src[3];
  dest[5] = src[2];
  dest[6] = src[1];
  dest[7] = src[0];

  return netlong;
}

#define ntohll(x) htonll(x)

struct WebSocketFrameHeader {
  uint8_t fin;
  uint8_t opcode;
  uint8_t mask;
  uint64_t payload_len;
  WebSocketFrameHeader(){};
  WebSocketFrameHeader(const std::string &data) {
    fin = (data[0] & 0x80) != 0;
    opcode = data[0] & 0x0f;
    mask = (data[1] & 0x80) != 0;
    payload_len = data[1] & 0x7f;
  }
  Task<bool> get_payload_len(FileStream &fconn) {
    if (payload_len > 127) {
      co_return false;
    }
    std::string data;
    if (payload_len == 126) {
      if ((data = co_await fconn.getn(2)).size() != 2) {
        co_return false;
      }
      payload_len = ntohs(*reinterpret_cast<uint16_t *>(data.data()));
      co_return true;
    }
    if (payload_len == 127) {
      if ((data = co_await fconn.getn(8)).size() != 8) {
        co_return false;
      }
      payload_len = ntohll(payload_len);
    }
    co_return true;
  }
};

WebSocketEndpoint::WebSocketEndpoint(WebSocket *manager, size_t id,
                                     FileStream &&conn)
    : m_manager(manager), m_id(id), m_conn(std::move(conn)) {
  uring_enqueue(loop, ping());
}

WebSocketEndpoint::~WebSocketEndpoint() {
  m_conn.~FileStream();
  m_manager->remove_endpoint(m_id);
}

// 读取一帧
Task<WebSocketEndpoint::Opcode> WebSocketEndpoint::recv(std::string &msg) {
  WebSocketFrameHeader header;
  while (true) {
    auto data1 = co_await m_conn.getn(2);
    header = {data1};
    if (header.opcode == CLOSE) {
      co_return CLOSE;
    }
    if (!co_await header.get_payload_len(m_conn)) {
      co_return CLOSE;
    };
    std::string mask_key;
    if (header.mask) {
      mask_key = co_await m_conn.getn(4);
    }
    std::string payload = co_await m_conn.getn(header.payload_len);
    if (header.opcode == PING) {
      msg = data1 + mask_key + payload;
      msg[0] = (msg[0] | 0x0F) & 0xFa;
      co_return PING;
    }
    if (header.opcode == PONG) {
      co_return PONG;
    }
    if (header.mask) {
      int idx = 0;
      for (auto &ch : payload) {
        ch = ch ^ mask_key[idx++ % 4];
      }
    }
    msg = payload;
    if (header.fin) {
      break;
    }
  }
  co_return (Opcode) header.opcode;
}

static std::string CompressPayload(const std::string &payload) {
  z_stream deflateStream;
  deflateStream.zalloc = Z_NULL;
  deflateStream.zfree = Z_NULL;
  deflateStream.opaque = Z_NULL;

  if (deflateInit(&deflateStream, Z_DEFAULT_COMPRESSION) != Z_OK) {
    throw std::runtime_error("Failed to initialize zlib for compression.");
  }

  deflateStream.avail_in = payload.size();
  deflateStream.next_in = (Bytef *)payload.c_str();

  char outbuffer[32768];
  deflateStream.avail_out = sizeof(outbuffer);
  deflateStream.next_out = (Bytef *)outbuffer;

  deflate(&deflateStream, Z_FINISH);
  deflateEnd(&deflateStream);

  return std::string(outbuffer, sizeof(outbuffer) - deflateStream.avail_out);
}

Task<bool> WebSocketEndpoint::send(const std::string &payload) {
  std::string data(2, 0);
  data[0] = (char)0x81;
  std::string compressed = payload;
  // std::string compressed = compress_payload(payload);
  if (compressed.size() < 126) {
    data[1] = (char)compressed.size();
  } else if (compressed.size() < 65536) {
    uint16_t compressed_len = compressed.size();
    compressed = htons(compressed_len);
    data[1] = 126;
    data += reinterpret_cast<const char *>(&compressed_len);
  } else {
    uint64_t compressed_len = compressed.size();
    compressed_len = htonll(compressed_len);
    data[1] = 127;
    data += reinterpret_cast<const char *>(&compressed_len);
  }
  std::cout << "ready to send" << std::endl;
  try {
    co_await m_conn.puts(data + compressed);
    co_await m_conn.flush();
  } catch (FileStream::ShutdownException &e) {
    co_return false;
  }
  co_return true;
}
Task<> WebSocketEndpoint::handle_ws() {
  while (true) {
    std::string msg;
    auto opcode = co_await recv(msg);
    if (opcode == CLOSE) {
      std::cout << "disconnect" << std::endl;
      co_return;
    }
    if (opcode == PING) {
      try {
        co_await m_conn.puts(msg);
        co_await m_conn.flush();
      } catch (FileStream::ShutdownException &e) {
        co_return;
      }
      continue;
    }
    if (opcode == PONG) {
      if (!m_ping) {
        co_return;
      }
      std::cout << "pong" << std::endl;
      m_pong = true;
      m_ping = false;
      continue;
    }

    auto res = co_await m_manager->m_msg_cb(*this, msg);
    if (!res) {
      co_return;
    }
  }
}

Task<> WebSocketEndpoint::ping() {
  while (true) {
    m_pong = false;
    if (m_ping) {
      co_return;
    }
    int res = co_await uring_timeout(loop, m_timeout, 0);
    if (res != 0) {
      co_return;
    }
    std::cout << res << std::endl;

    std::string data(2, 0);
    data[0] = (char)0x89;
    try {
      co_await m_conn.puts(data);
      co_await m_conn.flush();
    } catch (FileStream::ShutdownException &e) {
      m_conn.release();
      co_return;
    }

    m_ping = true;
    res = co_await uring_timeout(loop, m_timeout, 0);
    if (res != 0) {
      co_return;
    }
    if (!m_pong) {
      co_return;
    }
  }
}

WebSocket::WebSocket(const std::string &path) : m_path(path), m_index(0) {}

HttpResponse WebSocket::Handshake(const HttpRequest &req) {
  HttpResponse resp(101);
  auto it_ = req.get_header().find("sec-websocket-key");
  if (it_ == req.get_header().end()) {
    resp.m_status = 400;
    return resp;
  }
  resp.m_headers["connection"] = "Upgrade";
  resp.m_headers["upgrade"] = "websocket";
  std::string src = it_->second + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
  std::string target(28, 0);
  unsigned char sha1[20];
  SHA1((unsigned char *)src.data(), src.size(), sha1);
  EVP_EncodeBlock((unsigned char *)target.data(), sha1, 20);
  resp.m_headers["sec-websocket-accept"] = target;
  return resp;
}
void WebSocket::Update(HttpServer &server, const std::string &path) {
  auto it = server.m_handlers[Method::GET].find(path);
  bool exist = it != server.m_handlers[Method::GET].end();
  server.m_handlers[Method::GET][path] =
      [it, exist](const HttpRequest &req) -> Task<HttpResponse> {
    auto iupdate = req.get_header().find("upgrade");
    if (iupdate == req.get_header().end() || iupdate->second != "websocket") {
      if (exist) {
        co_return co_await it->second(req);
      }
      co_return HttpResponse(404);
    }
    co_return Handshake(req);
  };
}

WebSocketEndpoint::Ptr WebSocket::add_endpoint(FileStream &&conn) {
  WebSocketEndpoint::Ptr ws_endpoint =
      std::make_shared<WebSocketEndpoint>(this, m_index, std::move(conn));
  m_endpoints[m_index] = ws_endpoint;
  if (m_index == std::string::npos) {
    m_index = 0;
  } else {
    m_index++;
  }
  return ws_endpoint;
}
void WebSocket::remove_endpoint(size_t id) { m_endpoints.erase(id); }

void WebSocket::on_message(WsHandler task) { m_msg_cb = std::move(task); }

} // namespace co_async